package agent.componet.message.util;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.esotericsoftware.kryo.util.Pool;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

public class KryoUtils {
    // Kryo实例池（线程安全）
    private static final Pool<Kryo> kryoPool = new Pool<Kryo>(true, false, 8) {
        @Override
        protected Kryo create() {
            Kryo kryo = new Kryo();

            // 配置Kryo
            kryo.setReferences(false); // 禁用对象引用（提高速度）
            kryo.setRegistrationRequired(true); // 必须注册类（安全）
            kryo.setDefaultSerializer(com.esotericsoftware.kryo.serializers.FieldSerializer.class);

            // 注册常用JDK类
            kryo.register(Arrays.asList("").getClass());
            kryo.register(HashSet.class);

            return kryo;
        }
    };

    // 类注册表（确保所有序列化类已注册）
    private static final Set<Class<?>> registeredClasses = ConcurrentHashMap.newKeySet();

    /**
     * 注册需要序列化的类
     */
    public static void registerClass(Class<?> clazz) {
        if (registeredClasses.add(clazz)) {
            kryoPool.obtain().register(clazz);
            kryoPool.free(kryoPool.obtain());
        }
    }
    /**
     * 序列化对象（带Zstd压缩）
     */
    public static byte[] serialize(Object obj) {
        if (obj == null) return new byte[0];

        // 确保类已注册
        registerClass(obj.getClass());

        Kryo kryo = kryoPool.obtain();
        try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
             Output output = new Output(baos)) {

            kryo.writeClassAndObject(output, obj);
            output.flush();

            //
            return baos.toByteArray();

        } catch (Exception e) {
            throw new SerializationException("序列化失败", e);
        } finally {
            kryoPool.free(kryo);
        }
    }

    /**
     * 反序列化对象
     */
    @SuppressWarnings("unchecked")
    public static <T> T deserialize(byte[] data) {
        if (data == null || data.length == 0) return null;

        byte[] serialized = data;

        // 检查是否是压缩数据


        Kryo kryo = kryoPool.obtain();
        try (ByteArrayInputStream bais = new ByteArrayInputStream(serialized);
             Input input = new Input(bais)) {

            return (T) kryo.readClassAndObject(input);

        } catch (Exception e) {
            throw new SerializationException("反序列化失败", e);
        } finally {
            kryoPool.free(kryo);
        }
    }

    // 自定义序列化异常
    public static class SerializationException extends RuntimeException {
        public SerializationException(String message, Throwable cause) {
            super(message, cause);
        }

        public SerializationException(String message) {
            super(message);
        }
    }



}
